#include "usn.h"
#include <Windows.h>

#include "log.h"

anything::usn::UsnUtility::UsnUtility()
{
    
}

anything::usn::UsnUtility::~UsnUtility()
{

}

void anything::usn::UsnUtility::init_usn_information()
{
    // get logical drive string
    wchar_t drivers[drive_size];
    GetLogicalDriveStrings(drive_size,drivers);

    for(wchar_t* p=drivers;*p!='\0';p+=4)
    {
        if(DRIVE_FIXED==GetDriveType(p))
        {
            LOG_INFO(p);
            std::map<uint64,FileInfo>* files = new std::map<uint64,FileInfo>;
            _construct_drive_infonation(p,files);
            files_.push_back(files);
        }
    }

}

void anything::usn::UsnUtility::_construct_drive_infonation( const wchar_t* drive,std::map<uint64,FileInfo>* files )
{
    wchar_t file_system_name_buffer[16];
    ulong32 max_component_length = 0;
    int32 ret = GetVolumeInformation(drive,0,0,0,0,&max_component_length,file_system_name_buffer,16);
    if (0==ret)
    {
        return;
    }

    if(file_system_name_buffer[0]!=L'N' || file_system_name_buffer[1]!=L'T' || 
       file_system_name_buffer[2]!=L'F' || file_system_name_buffer[3]!=L'S')
    {
        return;
    }

    // drive should be like "\\.\XX:"
    std::wstring volumn_path = std::wstring(L"\\\\.\\")+*drive+L":";

    HANDLE volume_handle = NULL;
    do 
    {
        volume_handle = CreateFile(volumn_path.c_str(),
            GENERIC_READ | GENERIC_WRITE, 
            FILE_SHARE_READ | FILE_SHARE_WRITE,
            NULL,
            OPEN_EXISTING,
            0,
            NULL);

        if( volume_handle == INVALID_HANDLE_VALUE )
        {
            break;
        }

        LOG_INFO("Open the USN Journal file");
        // Open the USN Journal file
        ulong32 br;  
        CREATE_USN_JOURNAL_DATA cujd;  
        cujd.MaximumSize = 0;
        cujd.AllocationDelta = 0;
        int32 ret = DeviceIoControl(volume_handle,  
            FSCTL_CREATE_USN_JOURNAL,  
            &cujd,  
            sizeof(cujd),  
            NULL,  
            0,  
            &br,  
            NULL);  

        if(0==ret)
        {
            LOG_DEBUG("Open the USN Journal file error, the error code = "<<GetLastError());
            break;
        }


        // Queries for information on the current update sequence number (USN) change journal
        // its records, and its capacity
        LOG_INFO("Queries for information on the current USN");
        USN_JOURNAL_DATA usn_journal_data;
        ret = DeviceIoControl( volume_handle, FSCTL_QUERY_USN_JOURNAL, NULL, 0, &usn_journal_data, 
                               sizeof(usn_journal_data), &br, NULL );
        if(0==ret)
        {
            break;
            LOG_DEBUG("Queries for information on the current USN error, the error code = "<<GetLastError());
        }


        // Enumerates the update sequence number (USN) 
        LOG_INFO("Enumerates the update sequence number (USN)");
        MFT_ENUM_DATA mft_enum_data;  
        mft_enum_data.StartFileReferenceNumber = 0;  
        mft_enum_data.LowUsn = 0;
        mft_enum_data.HighUsn = usn_journal_data.NextUsn;
        
        //ulong32 lens = (max_component_length - 1) * sizeof(wchar_t) + sizeof(USN_RECORD);

        FileInfo file_info;
        
        char buffer[enum_usn_data_size];
        ulong32 usn_data_size = 0;
        wchar_t name[260];
        while (DeviceIoControl(volume_handle, FSCTL_ENUM_USN_DATA, &mft_enum_data, sizeof(mft_enum_data),
            buffer, enum_usn_data_size, &usn_data_size, NULL))
        {
            ulong32 bytes_returned = usn_data_size - sizeof(USN);
            USN_RECORD* usn_record = reinterpret_cast<USN_RECORD*>((char*)buffer+sizeof(USN));

            while( bytes_returned > 0 )
            {
                file_info.clear();
                memset(name,0,MAX_PATH*sizeof(wchar_t));
                file_info.file_ref_no = usn_record->FileReferenceNumber;
                file_info.parent_ref_no = usn_record->ParentFileReferenceNumber;
                file_info.file_attributes = usn_record->FileAttributes;
                memcpy(name, usn_record->FileName, usn_record->FileNameLength );
                file_info.name = name;
                (*files)[file_info.file_ref_no] = file_info;

                bytes_returned -= usn_record->RecordLength;
                usn_record = reinterpret_cast<USN_RECORD*>( (char*)usn_record + usn_record->RecordLength );
            }
            mft_enum_data.StartFileReferenceNumber=*(reinterpret_cast<USN*>(&buffer));
        }

        LOG_INFO("delete the USN Journal file");
        // delete the USN Journal file
        DELETE_USN_JOURNAL_DATA dujd;  
        dujd.UsnJournalID = usn_journal_data.UsnJournalID;  
        dujd.DeleteFlags = USN_DELETE_FLAG_DELETE;  

        int status = DeviceIoControl(volume_handle,  
            FSCTL_DELETE_USN_JOURNAL,  
            &dujd,  
            sizeof (dujd),  
            NULL,  
            0,  
            &br,  
            NULL);  

        if (0==ret){  
            break;
            LOG_DEBUG("delete the USN Journal file error, the error code = "<<GetLastError()); 
        }
    } while (false);
    CloseHandle(volume_handle);
}



void anything::usn::UsnUtility::FileInfo::clear()
{
    file_ref_no = 0;
    parent_ref_no = 0;
    file_attributes = 0;
    name.clear();
}
